import com.sun.jersey.api.client.ClientResponse;
import org.apache.atlas.ApplicationProperties;
import org.apache.atlas.AtlasClientV2;
import org.apache.atlas.AtlasServiceException;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
import org.apache.atlas.model.instance.AtlasEntityHeader;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.instance.EntityMutations;
import org.apache.atlas.utils.AuthenticationUtil;
import org.apache.atlas.hive.model.HiveDataTypes;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.commons.collections.PredicateUtils;
import org.apache.commons.configuration.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

import static org.apache.atlas.hive.hook.events.BaseHiveEvent.*;

/**
 * Created by sqy on 2018/12/11.
 */
public class AuditBridge {
    private static final Logger LOG = LoggerFactory.getLogger(AuditBridge.class);

    public static final String ATLAS_ENDPOINT = "atlas.rest.address";
    public static final String HIVE_CLUSTER_NAME  = "atlas.cluster.name";
    public static final String DEFAULT_CLUSTER_NAME  = "primary";

    private final String  clusterName;
    private final AtlasClientV2 atlasClientV2;
    private final static List<AuditRule> ruleList = JsonUtils.loadFromJson();

    public AuditBridge(Configuration atlasConf, AtlasClientV2 atlasClientV2) {
        this(atlasConf.getString(HIVE_CLUSTER_NAME,DEFAULT_CLUSTER_NAME),atlasClientV2);
    }

    public AuditBridge(String clusterName, AtlasClientV2 atlasClientV2) {
        this.clusterName = clusterName;
        this.atlasClientV2 = atlasClientV2;
    }

    public String getClusterName() {
        return clusterName;
    }

    public AtlasClientV2 getAtlasClientV2() {
        return atlasClientV2;
    }

    /*public static void main(String[] args) {

        final AtlasClientV2 atlasClientV2;
        try {
            Configuration atlasConf  = ApplicationProperties.get();
            String[]      atlasEndpoint    = atlasConf.getStringArray(ATLAS_ENDPOINT);

            if (!AuthenticationUtil.isKerberosAuthenticationEnabled()) {
                String[] basicAuthUsernamePassword = AuthenticationUtil.getBasicAuthenticationInput();

                atlasClientV2 = new AtlasClientV2(atlasEndpoint, basicAuthUsernamePassword);
            } else {
                UserGroupInformation ugi = UserGroupInformation.getCurrentUser();

                atlasClientV2 = new AtlasClientV2(ugi, ugi.getShortUserName(), atlasEndpoint);
            }

            AuditBridge auditBridge = new AuditBridge(atlasConf,atlasClientV2);
            auditBridge.importAudit();

        }  catch (Exception e) {
            e.printStackTrace();
        }

    }*/

    private void importAudit() throws Exception {
        if (CollectionUtils.isNotEmpty(ruleList)){
            for (AuditRule rule : ruleList) {
                registerPurpose(rule);
                registerTask(rule);
                registerUsedData(rule);
            }
        }
    }

    private  AtlasEntityWithExtInfo registerPurpose(AuditRule rule) throws Exception {
        AtlasEntityWithExtInfo ret = null;
        ret = findPurpose(rule.getPurpose());

        if (ret == null){
            ret = registerInstance(new AtlasEntityWithExtInfo(toPurposeEntity(rule)));
        }else {

            ret.setEntity(toPurposeEntity(rule,ret.getEntity()));

            updateInstance(ret);
        }

        return ret;
    }

    private AtlasEntityWithExtInfo registerTask(AuditRule rule) throws Exception {
        AtlasEntityWithExtInfo ret = null;

        Set<String> tasks = rule.getTask();
        if (CollectionUtils.isNotEmpty(tasks)){
            for (String task : tasks) {
                ret = findTask(task);
                if (ret == null){
                    ret = registerInstance(new AtlasEntityWithExtInfo(toTaskEntity(task,rule)));
                }else {

                    ret.setEntity(toTaskEntity(task,rule,ret.getEntity()));

                    updateInstance(ret);
                }
            }
        }

        return ret;
    }

    private AtlasEntityWithExtInfo registerUsedData(AuditRule rule) throws Exception {
        AtlasEntityWithExtInfo ret = null;

        Set<String> usedDatas = rule.getUsedData();
        if (CollectionUtils.isNotEmpty(usedDatas)){
            for (String usedData : usedDatas) {
                if (checkPunctuation(usedData)){
                    String[] strings = usedData.split(".");
                    String col = strings[1];
                    ret = findCol(col);

                    if (ret == null){
                        ret = registerInstance(new AtlasEntityWithExtInfo(toColumnEntity(col,rule)));
                    }else{
                        ret.setEntity(toColumnEntity(col,rule,ret.getEntity()));
                        updateInstance(ret);
                    }
                }else {
                    ret = findTable(usedData);
                    if (ret == null){
                        ret = registerInstance(new AtlasEntityWithExtInfo(toTableEntity(usedData,rule)));
                    }else{
                        ret.setEntity(toTableEntity(usedData,rule,ret.getEntity()));
                        updateInstance(ret);
                    }
                }

            }
        }

        return ret;
    }


    private  AtlasEntity toPurposeEntity(AuditRule rule) {
        return toPurposeEntity(rule,null);
    }

    private AtlasEntity toPurposeEntity(AuditRule rule, AtlasEntity entity) {
        if(entity == null){
            entity = new AtlasEntity(HiveDataTypes.HIVE_INDEX.getName());
        }

        entity.setAttribute(ATTRIBUTE_NAME,rule.getPurpose());
        entity.setAttribute(ATTRIBUTE_DESCRIPTION,rule.getDescription());
        entity.setAttribute(ATTRIBUTE_CREATE_TIME,rule.getCreateTime());
        entity.setAttribute(ATTRIBUTE_USER_NAME,rule.getCreateBy());

        return entity;
    }

    private  AtlasEntity toTaskEntity(String task,AuditRule rule) {
        return toTaskEntity(task,rule,null);
    }

    private AtlasEntity toTaskEntity(String task,AuditRule rule, AtlasEntity entity) {
        if(entity == null){
            entity = new AtlasEntity(HiveDataTypes.HIVE_TYPE.getName());
        }
        entity.setAttribute(ATTRIBUTE_NAME,task);
        entity.setAttribute(ATTRIBUTE_DESCRIPTION,rule.getDescription());
        entity.setAttribute(ATTRIBUTE_CREATE_TIME,rule.getCreateTime());
        entity.setAttribute(ATTRIBUTE_USER_NAME,rule.getCreateBy());

        return entity;

    }

    private AtlasEntity toTableEntity(String table, AuditRule rule, AtlasEntity entity) {
        if(entity == null){
            entity = new AtlasEntity(HiveDataTypes.HIVE_TABLE.getName());
        }
        entity.setAttribute(ATTRIBUTE_NAME,table);
        entity.setAttribute(ATTRIBUTE_INPUTS,rule.getPurpose());
        entity.setAttribute(ATTRIBUTE_OUTPUTS,rule.getTask());
        entity.setAttribute(ATTRIBUTE_DESCRIPTION,rule.getDescription());
        entity.setAttribute(ATTRIBUTE_CREATE_TIME,rule.getCreateTime());
        entity.setAttribute(ATTRIBUTE_USER_NAME,rule.getCreateBy());

        return entity;
    }

    private AtlasEntity toTableEntity(String table, AuditRule rule) {
        return toTableEntity(table,rule,null);
    }

    private AtlasEntity toColumnEntity(String col, AuditRule rule) {
        return toColumnEntity(col,rule,null);
    }

    private AtlasEntity toColumnEntity(String col, AuditRule rule, AtlasEntity entity) {
        if(entity == null){
            entity = new AtlasEntity(HiveDataTypes.HIVE_COLUMN.getName());
        }
        entity.setAttribute(ATTRIBUTE_NAME,col);
        entity.setAttribute(ATTRIBUTE_INPUTS,rule.getPurpose());
        entity.setAttribute(ATTRIBUTE_OUTPUTS,rule.getTask());
        entity.setAttribute(ATTRIBUTE_DESCRIPTION,rule.getDescription());
        entity.setAttribute(ATTRIBUTE_CREATE_TIME,rule.getCreateTime());
        entity.setAttribute(ATTRIBUTE_USER_NAME,rule.getCreateBy());

        return entity;
    }


    private AtlasEntityWithExtInfo findPurpose(String purpose) {
        String typeName = HiveDataTypes.HIVE_INDEX.getName();

        return findPurposeEntity(typeName,purpose);
    }

    private AtlasEntityWithExtInfo findPurposeEntity(final String typeName, final String purpose) {
        AtlasEntityWithExtInfo ret = null;
        try {
            ret = atlasClientV2.getEntityByAttribute(typeName, Collections.singletonMap(ATTRIBUTE_NAME,purpose));
        } catch (AtlasServiceException e) {
            if (e.getStatus() == ClientResponse.Status.NOT_FOUND){
                return null;
            }
        }
       // clearRelationshipAttributes(ret);
        return ret;
    }

    private AtlasEntityWithExtInfo findTask(String task) {
        String typeName = HiveDataTypes.HIVE_TYPE.getName();

        return findTaskEntity(typeName,task);

    }

    private AtlasEntityWithExtInfo findTaskEntity(String typeName, String task) {
        AtlasEntityWithExtInfo ret = null;
        try {
            ret = atlasClientV2.getEntityByAttribute(typeName, Collections.singletonMap(ATTRIBUTE_NAME,task));
        } catch (AtlasServiceException e) {
            if (e.getStatus() == ClientResponse.Status.NOT_FOUND){
                return null;
            }
        }
        // clearRelationshipAttributes(ret);
        return ret;
    }

    private AtlasEntityWithExtInfo findTable(String table) {
        String typeName = HiveDataTypes.HIVE_TABLE.getName();
        return findTableEntity(typeName,table);
    }

    private AtlasEntityWithExtInfo findTableEntity(String typeName, String table) {
        AtlasEntityWithExtInfo ret = null;
        try {
            ret = atlasClientV2.getEntityByAttribute(typeName, Collections.singletonMap(ATTRIBUTE_NAME,table));
        } catch (AtlasServiceException e) {
            if (e.getStatus() == ClientResponse.Status.NOT_FOUND){
                return null;
            }
        }
        // clearRelationshipAttributes(ret);
        return ret;
    }

    private AtlasEntityWithExtInfo findCol(String col) {
        String typeName = HiveDataTypes.HIVE_COLUMN.getName();
        return findColEntity(typeName,col);
    }

    private AtlasEntityWithExtInfo findColEntity(final String typeName, final String col) {
        AtlasEntityWithExtInfo ret = null;
        try {
            ret = atlasClientV2.getEntityByAttribute(typeName, Collections.singletonMap(ATTRIBUTE_NAME,col));
        } catch (AtlasServiceException e) {
            if (e.getStatus() == ClientResponse.Status.NOT_FOUND){
                return null;
            }
        }
        // clearRelationshipAttributes(ret);
        return ret;
    }

    private AtlasEntityWithExtInfo registerInstance(AtlasEntityWithExtInfo entity) throws Exception {

        AtlasEntityWithExtInfo  ret             = null;
        EntityMutationResponse response        = atlasClientV2.createEntity(entity);
        List<AtlasEntityHeader> createdEntities = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

        if (CollectionUtils.isNotEmpty(createdEntities)) {
            for (AtlasEntityHeader createdEntity : createdEntities) {
                if (ret == null) {
                    ret = atlasClientV2.getEntityByGuid(createdEntity.getGuid());

                    LOG.info("Created {} entity: name={}, guid={}", ret.getEntity().getTypeName(), ret.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), ret.getEntity().getGuid());
                } else if (ret.getEntity(createdEntity.getGuid()) == null) {
                    AtlasEntityWithExtInfo newEntity = atlasClientV2.getEntityByGuid(createdEntity.getGuid());

                    ret.addReferredEntity(newEntity.getEntity());

                    if (MapUtils.isNotEmpty(newEntity.getReferredEntities())) {
                        for (Map.Entry<String, AtlasEntity> entry : newEntity.getReferredEntities().entrySet()) {
                            ret.addReferredEntity(entry.getKey(), entry.getValue());
                        }
                    }

                    LOG.info("Created {} entity: name={}, guid={}", newEntity.getEntity().getTypeName(), newEntity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), newEntity.getEntity().getGuid());
                }
            }
        }

       // clearRelationshipAttributes(ret);

        return ret;
    }
    private void updateInstance(AtlasEntityWithExtInfo entity) throws AtlasServiceException {

        atlasClientV2.updateEntity(entity);

        LOG.info("Updated {} entity: name={}, guid={}", entity.getEntity().getTypeName(), entity.getEntity().getAttribute(ATTRIBUTE_QUALIFIED_NAME), entity.getEntity().getGuid());
    }

    private void clearRelationshipAttributes(AtlasEntityWithExtInfo entity) {

        if (entity != null) {
            clearRelationshipAttributes(entity.getEntity());

            if (entity.getReferredEntities() != null) {
                clearRelationshipAttributes(entity.getReferredEntities().values());
            }
        }
    }

    private void clearRelationshipAttributes(Collection<AtlasEntity> entities) {

        if (entities != null) {
            for (AtlasEntity entity : entities) {
                clearRelationshipAttributes(entity);
            }
        }
    }

    private void clearRelationshipAttributes(AtlasEntity entity) {
        if (entity != null && entity.getRelationshipAttributes() != null) {
            entity.getRelationshipAttributes().clear();
        }
    }

    //检查使用的数据类型是整张表还是某些列
    private boolean checkPunctuation(String usedData){
        boolean flag = false;
        String tmp = usedData;
        tmp = tmp.replace("\\p{P}","");

        if (tmp.length() != usedData.length()){
            flag = true;
        }
        return flag;
    }

}
